#!/bin/bash

# Default values
OUTPUT_DIR=""
SIZE_MB=""
FILE_COUNT=""

# --- Helper Functions ---
usage() {
    echo "Usage: $0 --output <dir_name> --size <size_MB_per_file> --count <num_files>"
    echo "  --output: Directory name to store the generated task."
    echo "  --size: Size of *each* file in MB (e.g., 1, 100, 1000)."
    echo "  --count: Number of files to generate (e.g., 1000, 10, 1)."
    exit 1
}

print_error() {
    echo "Error: $1" >&2
    usage
}

# --- Argument Parsing ---
while [[ $# -gt 0 ]]; do
    key="$1"
    case $key in
        --output)
        OUTPUT_DIR="$2"
        shift; shift
        ;;
        --size)
        SIZE_MB="$2"
        shift; shift
        ;;
        --count)
        FILE_COUNT="$2"
        shift; shift
        ;;
        *) # unknown option
        print_error "Unknown option $1"
        ;;
    esac
done

# --- Input Validation ---
if [[ -z "$OUTPUT_DIR" || -z "$SIZE_MB" || -z "$FILE_COUNT" ]]; then
    print_error "Missing required arguments (--output, --size, --count)."
fi

if ! [[ "$SIZE_MB" =~ ^[0-9]+$ ]] || [[ "$SIZE_MB" -le 0 ]]; then
    print_error "Size must be a positive integer."
fi

if ! [[ "$FILE_COUNT" =~ ^[0-9]+$ ]] || [[ "$FILE_COUNT" -le 0 ]]; then
    print_error "Count must be a positive integer."
fi


if [[ -e "$OUTPUT_DIR" ]]; then
    echo "Warning: Output directory '$OUTPUT_DIR' already exists. Files might be overwritten."
else
    mkdir -p "$OUTPUT_DIR" || { echo "Error: Failed to create output directory '$OUTPUT_DIR'." >&2; exit 1; }
fi

# --- Task Generation ---
echo "Generating task in directory '$OUTPUT_DIR'..."
echo "Creating $FILE_COUNT file(s), each ${SIZE_MB}MB..."

block_size="1M" # Use 1MB blocks for dd (actually MiB, 1024*1024)
total_size_mib=$((SIZE_MB * FILE_COUNT))
total_size_bytes=$((total_size_mib * 1024 * 1024)) # Calculate actual total bytes

echo "Total size: ${total_size_mib} MiB (${total_size_bytes} bytes)"

for i in $(seq 1 "$FILE_COUNT"); do
    # Use a consistent naming scheme, even for a single file
    output_file="$OUTPUT_DIR/file_$i.dat"
    # Use status=none for less verbose output during loop
    dd if=/dev/zero of="$output_file" bs="$block_size" count="$SIZE_MB" status=none || { echo "Error: dd command failed for file $i." >&2; exit 1; }
    # Use /dev/urandom for non-zero data, but it's much slower:
    # dd if=/dev/urandom of="$output_file" bs="$block_size" count="$SIZE_MB" status=none || exit 1

    # Optional progress indicator
    if [[ $(($i % 100)) -eq 0 || $i -eq $FILE_COUNT ]]; then
         printf "\rProgress: %d/%d" "$i" "$FILE_COUNT"
    fi
done
echo # Newline after progress

echo "Task generation complete in '$OUTPUT_DIR'."
exit 0
